from scipy.optimize import minimize, curve_fit, root
from typing import Callable, Tuple, Optional, Union, List
import numpy as np
import matplotlib.pyplot as plt

def optimize_min_value(
    fun: Callable, 
    origin_guess: np.ndarray, 
    method: str = 'BFGS',
    gtol: float = 1e-6,
    maxiter: int = int(1e6)
) -> Tuple[np.ndarray, float]:
    """使用优化算法寻找函数的最小值点
    
    参数:
        fun: 目标函数，接受numpy数组作为输入，返回一个标量值
        origin_guess: 初始猜测点
        method: 优化算法，默认使用'BFGS'
        gtol: 梯度收敛阈值
        maxiter: 最大迭代次数
    
    返回:
        Tuple[np.ndarray, float]: (最小值点坐标, 最小值)
    """
    result = minimize(
        fun=fun,
        x0=origin_guess,
        method=method,
        options={'disp': True, 'gtol': gtol, 'maxiter': maxiter},
    )
    print(f"最小值点：{result.x}")
    print(f"最小值：{result.fun}")
    print(f"迭代次数：{result.nit}")
    return result.x, result.fun

def fit_rational_function(
    x: np.ndarray,
    y: np.ndarray,
    p0: Optional[list[float]] = None
) -> Tuple[np.ndarray, float]:
    """拟合有理分式函数 y = (a + b*x)/(c*x)
    
    参数:
        x: 自变量数据
        y: 因变量数据
        p0: 初始参数猜测 [a, b, c]
    
    返回:
        Tuple[np.ndarray, float]: (最优参数 [a, b, c], R方值)
    """
    def rational_func(x, a, b, c):
        return (a + b*x)/(c*x)
    
    # 如果没有提供初始猜测，使用默认值
    if p0 is None:
        p0 = [1000, -0.1, 0.1]
    
    # 拟合
    popt, pcov = curve_fit(rational_func, x, y, p0=p0)
    
    # 计算R方值
    y_pred = rational_func(x, *popt)
    r2 = 1 - np.sum((y - y_pred)**2) / np.sum((y - np.mean(y))**2)
    
    return popt, r2

def solve_equations(
    func: Callable[[np.ndarray], np.ndarray],
    x0: Union[List[float], np.ndarray],
    method: str = 'hybr',
    tol: float = 1e-6,
    max_iter: int = 1000
) -> Tuple[np.ndarray, bool]:
    """求解非线性方程组 f(x) = 0
    
    参数:
        func: 方程组函数，接受一个向量x并返回一个等长的向量f(x)
        x0: 初始猜测值
        method: 求解方法，可选：'hybr', 'lm', 'broyden1', 'broyden2', 'anderson'等
        tol: 收敛容差
        max_iter: 最大迭代次数
    
    返回:
        Tuple[np.ndarray, bool]: (解向量, 是否成功)
        
    示例:
        def equations(vars):
            x, y = vars
            return [x**2 + y**2 - 1,     # 第一个方程：x² + y² = 1
                   x - y]                 # 第二个方程：x = y
        
        solution, success = solve_equations(equations, [1.0, 1.0])
    """
    result = root(
        func,
        x0=np.array(x0),
        method=method,
        options={'xtol': tol, 'maxiter': max_iter}
    )
    
    if result.success:
        print(f"方程组求解成功:")
        print(f"解向量: {result.x}")
        print(f"迭代次数: {result.nfev}")
        print(f"收敛误差: {np.max(np.abs(result.fun))}")
    else:
        print("方程组求解失败")
        print(f"错误信息: {result.message}")
    
    return result.x, result.success
if __name__ == "__main__":
    # 测试最小值寻找
    print("测试1: 寻找函数 f(x,y) = x^2 + y^2 的最小值")
    def test_func(x):
        return x[0]**2 + x[1]**2
    
    min_point, min_value = optimize_min_value(test_func, np.array([100, -100]))
    print(f"预期最小值点: [0, 0]")
    print(f"预期最小值: 0\n")

    # 测试3: 求解非线性方程组
    print("\n测试3: 求解非线性方程组")
    def test_equations(vars):
        x, y = vars
        return np.array([
            x**2 + y**2 - 1,  # 第一个方程：x² + y² = 1 (圆方程)
            x - y             # 第二个方程：x = y (直线方程)
        ])
    
    # 求解方程组（应该在x=y的直线与单位圆的交点处）
    solution, success = solve_equations(test_equations, [1.0, 1.0])
    
    if success:
        # 验证解的准确性
        x, y = solution
        print("\n验证解:")
        print(f"x² + y² = {x**2 + y**2:.6f} (应接近1)")
        print(f"x - y = {x - y:.6f} (应接近0)")
        
        # 在新图中可视化解
        plt.figure(figsize=(8, 8))
        
        # 画圆
        theta = np.linspace(0, 2*np.pi, 100)
        circle_x = np.cos(theta)
        circle_y = np.sin(theta)
        plt.plot(circle_x, circle_y, 'b-', label='x² + y² = 1')
        
        # 画直线
        line_x = np.linspace(-1.5, 1.5, 100)
        plt.plot(line_x, line_x, 'r-', label='x = y')
        
        # 标出解点
        plt.scatter([x], [y], color='green', s=100, label='Solution')
        
        plt.title('Visualization of System of Equations')
        plt.xlabel('x')
        plt.ylabel('y')
        plt.axis('equal')
        plt.grid(True)
        plt.legend()
        plt.show()
    
    # 测试有理函数拟合
    # print("测试2: 拟合有理函数")
    # # 生成测试数据
    # x_data = np.array([124853, 139164, 159778, 162974])
    # y_data = np.array([8.838163, 8.021931, 7.103200, 6.982578])
    
    # # 拟合
    # params, r2 = fit_rational_function(x_data, y_data)
    # a, b, c = params
    # print(f"拟合参数:")
    # print(f"a = {a:.6f}")
    # print(f"b = {b:.6f}")
    # print(f"c = {c:.6f}")
    # print(f"R² = {r2:.6f}")
    
    # # 绘制拟合结果
    # plt.figure(figsize=(10, 6))
    # plt.scatter(x_data, y_data, color='blue', label='Data points')
    
    # x_fit = np.linspace(min(x_data), max(x_data), 100)
    # y_fit = (a + b*x_fit)/(c*x_fit)
    # plt.plot(x_fit, y_fit, 'r-', label=f'Fit: y=({a:.1f}+{b:.6f}x)/({c:.6f}x)')
    
    # plt.title('Rational Function Fitting')
    # plt.xlabel('x')
    # plt.ylabel('y')
    # plt.legend()
    # plt.grid(True)
    # plt.show()

